
#include "RSADecrypt.h"

/** Construct the class with no bucket brigades. This means that this
 * is to be used for block and not pipe decryption.
 * 
 * @param rsaprvkey The RSA private key to use for decryption. */
RSADecrypt::RSADecrypt(const RSAPrivateKey &rsaprvkey)
{
   rsaprvkey.getSSLKeyComponents(&key);
   input = NULL;
   output = NULL;
   previous_block = NULL;
}

/** Construct the class with bucket brigades. This means that this can
 * be used for both block and pipe decryption.
 * 
 * @param iinput The bucket brigade for input.
 * 
 * @param ioutput The bucket brigade for output.
 * 
 * @param rsaprvkey The RSA private key to use for decryption. */
RSADecrypt::RSADecrypt(APRBucketBrigade &iinput, APRBucketBrigade &ioutput, const RSAPrivateKey &rsaprvkey)
{
   rsaprvkey.getSSLKeyComponents(&key);
   input = &iinput;
   output = &ioutput;
   previous_block = NULL;
}

/** Destructor for the RSA private key decryption class. */
RSADecrypt::~RSADecrypt()
{
   RSA_free(key);
}

/** Decrypt data block and return plaintext data. The decrypted data
 * returned must be freed by the caller.
 * 
 * @param codedblock A block of encrypted data to decrypt.
 * 
 * @param codedsize The length of the encrypted data block.
 * 
 * @param[out] plainsize The length of the plaintext block of data
 * returned. The data returned must be freed by the caller. */
char *RSADecrypt::processDataBlock(const char *codedblock, apr_size_t codedsize, apr_size_t *plainsize)
{
   if (codedblock == NULL)
   {
      errprint("tried to decrypt NULL data block");
      return NULL;
   }
   if (key == NULL)
   {
      errprint("key is NULL - decryption failed");
      return NULL;
   }
   // For the PKCS v1.5 padding to be included too, apparently.
   apr_size_t codeddata_blocksze = RSA_size(key);
   apr_size_t plaindata_blocksze = RSA_size(key) - RSA_PADDING_SIZE;
   int noofcodeddatablocks = ((codeddata_blocksze + codedsize - 1) / (codeddata_blocksze));
   apr_size_t t_maxplainsize = noofcodeddatablocks * codeddata_blocksze;
   char *returnblock = (char *)malloc(sizeof(char)*t_maxplainsize);
   previous_block = (char *)calloc(codeddata_blocksze, sizeof(char));
   char *manipdata = (char *)malloc(sizeof(char)*codeddata_blocksze);
   int j = 0;
   for (int i = 0; i < codedsize; i += codeddata_blocksze)
   {
      int tcodedblock = ((codedsize - i) < codeddata_blocksze) ? (codedsize - i) : codeddata_blocksze;
      int thisblock = RSA_private_decrypt(tcodedblock,
         (const unsigned char *)(&(codedblock[i])),
         (unsigned char *)(manipdata), key, RSA_PKCS1_PADDING);
      if (thisblock < 0)
      {
         SSLCheckError();
         free(returnblock);
         free(manipdata);
         free(previous_block);
         return NULL;
      } else
      {
         for (int k = 0; k < thisblock; k++)
            manipdata[k] ^= previous_block[k];
         memcpy(&(returnblock[j]), manipdata, sizeof(char)*thisblock);
         j += thisblock;
         memcpy(previous_block, &(codedblock[i]), sizeof(char)*tcodedblock);
      }
   }
   free(manipdata);
   if (plainsize) *plainsize = j;
   free(previous_block);
   previous_block = NULL;
   return returnblock;
}

/** Begin the process of decrypting a pipe of data.
 * 
 * @return Zero on success, non-zero on failure. */
int RSADecrypt::processPipeBegin()
{
   if (input == NULL)
   {
      errprint("no input pipe (is NULL)");
      return 1;
   }
   if (output == NULL)
   {
      errprint("no output pipe (is NULL)");
      return 1;
   }
   if (key == NULL)
   {
      errprint("key is NULL - encryption failed");
      return 1;
   }
   // In the RSA decryption there is no state, but we should allocate
   // the previous block
   previous_block = (char *)calloc(RSA_size(key), sizeof(char));
}

/** Process any outstanding encrypted data on the input pipe, and
 * place any processed decrypted data into the output pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int RSADecrypt::processPipeStep()
{
   // Padding
   apr_size_t codeddata_blocksize = RSA_size(key);
   apr_size_t plaindata_blocksize = RSA_size(key) - RSA_PADDING_SIZE;
   char *itemp_block = (char *)malloc(sizeof(char)*codeddata_blocksize);
   char *otemp_block = (char *)malloc(sizeof(char)*plaindata_blocksize);
   
   while (1)
   {
      apr_size_t readbytes = input->readFromStream(itemp_block, codeddata_blocksize);
      // this isn't a problem, this should fail during normal
      // operation when the buffer runs out of data.
      if (readbytes != codeddata_blocksize) break;
      int thisblock = RSA_private_decrypt(codeddata_blocksize,
         (const unsigned char *)itemp_block, (unsigned char *)otemp_block, key, RSA_PKCS1_PADDING);
      for (int k = 0; k < thisblock; k++)
         otemp_block[k] ^= previous_block[k];
      apr_size_t writebytes = output->writeToStream(otemp_block, thisblock);
      if (writebytes != thisblock)
      {
         errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)thisblock);
         free(itemp_block);
         free(otemp_block);
         return 1;
      }
      memcpy(previous_block, itemp_block, codeddata_blocksize);
      input->removeFromStream(codeddata_blocksize);
   }
   free(itemp_block);
   free(otemp_block);
   return 0;
}

/** Process the pipe step and record the movement of data to determine
 * when the process is complete.
 * 
 * @param[out] movementin Records data input in bytes.
 * 
 * @param[out] movementout Records data output in bytes. */
int RSADecrypt::processPipeStep(int *movementin, int *movementout)
{
   apr_size_t before_input = input->realLength();
   apr_size_t before_output = output->realLength();
   int retval = processPipeStep();
   apr_size_t after_input = input->realLength();
   apr_size_t after_output = output->realLength();
   if (movementin)
      *movementin = before_input - after_input;
   if (movementout)
      *movementout = after_output - before_output;
   return retval;
}

/** Process any outstanding plaintext data in the pipe and finalise
 * the encryption process.
 * 
 * @return Zero on success, non-zero on failure. */
int RSADecrypt::processPipeEnd()
{
   processPipeStep();
   // Padding
   apr_size_t codeddata_blocksize = input->realLength();
   if (codeddata_blocksize > (RSA_size(key)))
   {
      errprint("too much data remaining in bucket brigade (is it still being written to?)");
      return 1;
   }
   char *itemp_block = (char *)malloc(sizeof(char)*codeddata_blocksize);
   // should be smaller than the coded data block!
   char *otemp_block = (char *)malloc(sizeof(char)*codeddata_blocksize);
   apr_size_t readbytes = input->readFromStream(itemp_block, codeddata_blocksize);
   // this isn't a problem, this should fail during normal operation
   // when the buffer runs out of data.
   if (readbytes != codeddata_blocksize)
   {
      errprint("failed to read data from plain bucket brigade (read %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), buffer unfinished", readbytes, codeddata_blocksize);
      free(itemp_block);
      free(otemp_block);
      return 1;
   }

   if (codeddata_blocksize > 0)
   {
      int thisblock = RSA_public_encrypt(codeddata_blocksize,
         (const unsigned char *)itemp_block, (unsigned char *)otemp_block, key, RSA_PKCS1_PADDING);
      for (int k = 0; k < thisblock; k++)
         otemp_block[k] ^= previous_block[k];
      apr_size_t writebytes = output->writeToStream(otemp_block, thisblock);
      if (writebytes != thisblock)
      {
         errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)thisblock);
         free(itemp_block);
         free(otemp_block);
         return 1;
      }
      input->removeFromStream(codeddata_blocksize);
   }
   free(itemp_block);
   free(otemp_block);
   free(previous_block);
   previous_block = NULL;
   return 0;
}
